package com.baomidou.mybatisplus.samples.tenant.config;

import com.baomidou.mybatisplus.core.toolkit.ArrayUtils;
import com.baomidou.mybatisplus.samples.tenant.context.SecurityContextHandler;
import org.springframework.aop.interceptor.AsyncUncaughtExceptionHandler;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.TaskDecorator;
import org.springframework.scheduling.annotation.AsyncConfigurerSupport;
import org.springframework.scheduling.annotation.EnableAsync;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.Arrays;
import java.util.concurrent.Executor;
import java.util.concurrent.ThreadPoolExecutor;

/**
 * 注解@Async指定自定义线程池, 保证线程间认证数据，账号数据不丢失
 *
 * @Auther: 乔帅 ymxz34787409@126.com
 * @Date: 2022/2/21 14:35
 */
@Configuration
@EnableAsync
public class AsyncConfiguration extends AsyncConfigurerSupport {


    @Bean
    public ThreadPoolTaskExecutor executor() {
        final ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor();
        // 核心线程数
        taskExecutor.setCorePoolSize(10);
        // 最大线程数
        taskExecutor.setMaxPoolSize(50);
        // 队列大小
        taskExecutor.setQueueCapacity(10);
        // 队列满后拒绝策略 - > CallerRunsPolicy由调用线程处理该任务
        taskExecutor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
        // 线程名前缀
        String threadNamePrefix = "ynx:thread-";
        taskExecutor.setThreadNamePrefix(threadNamePrefix);
        // 该方法就是这里的关键，用来设置线程池关闭的时候等待所有任务都完成再继续销毁其他的Bean，
        taskExecutor.setWaitForTasksToCompleteOnShutdown(true);
        // 使用自定义的跨线程的请求级别线程工厂类
        int awaitTerminationSeconds = 60;
        // 线程池中任务的等待时间，如果超过这个时间还没有销毁就 强制销毁，以确保应用最后能够被关闭，而不是阻塞住
        taskExecutor.setAwaitTerminationSeconds(awaitTerminationSeconds);
        // 添加装饰器模式 增强了Runnable接口, 透传上下文等参数
        taskExecutor.setTaskDecorator(new PlusDecorator());
        // 初始化线程池
        taskExecutor.initialize();
        return taskExecutor;

    }


    /**
     * 异步执行指定自定义线程池保证权限信息的传递
     */
    @Override
    public Executor getAsyncExecutor() {
        return executor();
    }

    /**
     * 线程异常处理
     */
    @Override
    public AsyncUncaughtExceptionHandler getAsyncUncaughtExceptionHandler() {
        return (throwable, method, objects) -> {
            throwable.printStackTrace();
            StringBuilder sb = new StringBuilder();
            sb.append("异步执行异常 Exception message - ").append(throwable.getMessage())
                .append(", 异步方式 Method name - ").append(method.getName());
            if (ArrayUtils.isNotEmpty(objects)) {
                sb.append(",异步参数 Parameter value - ").append(Arrays.toString(objects));
            }
            throw new RuntimeException(sb.toString());
        };
    }


    /**
     * 线程装饰，增强Runnable。 可以传递认证信息和请求上下文属性
     */
    class PlusDecorator implements TaskDecorator {

        /**
         * 增强runable方法
         */
        @Override
        public Runnable decorate(Runnable runnable) {
            try {
                /*
                主线程获取需要透传给子线程的变量,
                1. 请求上下文RequestAttributes(含请求头等信息)
                2. 登陆的用户信息 loginUser
                 */
//                final RequestAttributes attr = RequestContextHolder.getRequestAttributes();
                final Boolean ignore = SecurityContextHandler.getIgnore();
                return () -> {
                    // 这里是 Runable接口run方法，是子线程执行的
                    try {
                        // 子线程赋值传过来的变量
//                        RequestContextHolder.setRequestAttributes(attr);
                        // 设置主线程传过来的用户数据
//                        LoginHelper.setUserThreadLocal(loginUser);
                        // 设置主线程传过来的是否忽略租户条件数据
                        SecurityContextHandler.setIgnore(ignore);
                        runnable.run();
                    } catch (Exception e) {
                        e.printStackTrace();
                    } finally {
                        // 清空attributes属性
//                        RequestContextHolder.resetRequestAttributes();
                    }
                };

            } catch (IllegalStateException e) {
                return runnable;
            }
        }
    }
}
